﻿//
// Copyright (c) 2010 Alex Forster.
// 
// Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
// documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
// permit persons to whom the Software is furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in all copies or substantial portions of
// the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
// TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
// 
// ZeroMQ is a trademark of iMatix Corporation.
// 

using System;

using System.Text;
using System.Threading;
using System.ComponentModel;
using System.Runtime.InteropServices;
using System.Collections.Generic;

namespace ZeroMQ
{
	public delegate void ZmqSendEvent( Byte[] data, Boolean? multipart );

	public delegate void ZmqReceiveEvent();

	[EditorBrowsable( EditorBrowsableState.Never )]
	public abstract class ZmqSocket: MarshalByRefObject, IDisposable
	{
		private readonly IntPtr m_send_msg = Marshal.AllocHGlobal( IntPtr.Size + 32 );

		private readonly IntPtr m_receive_msg = Marshal.AllocHGlobal( IntPtr.Size + 32 );

		private static LinkedList<ZmqSocket> s_socket_list;

		private static Thread s_poller;

		private static IntPtr s_context;

		public readonly IntPtr Context;

		public readonly IntPtr Socket;

		protected ZmqSocket( Int32 type, IntPtr? context )
		{
			if( s_context == IntPtr.Zero )
			{
				s_socket_list = new LinkedList<ZmqSocket>();

				s_poller = new Thread( ZmqSocket.Schedule )
				{
					IsBackground = true
				};

				s_poller.Start();

				s_context = Interop.zmq_init( Environment.ProcessorCount );

				if( s_context == IntPtr.Zero )
				{
					throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
				}
			}

			if( context != null && context.Value != IntPtr.Zero )
			{
				this.Context = context.Value;
			}
			else
			{
				this.Context = s_context;
			}

			this.Socket = Interop.zmq_socket( this.Context, type );

			if( this.Socket == IntPtr.Zero )
			{
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			lock( s_socket_list ) // note- expensive to raipdly create or destroy sockets
			{
				s_socket_list.AddLast( this );
			}
		}

		private String m_version;

		public String Version
		{
			get
			{
				if( String.IsNullOrEmpty( m_version ) )
				{
					var p_major = Marshal.AllocHGlobal( Marshal.SizeOf( typeof( Int32 ) ) );
					var p_minor = Marshal.AllocHGlobal( Marshal.SizeOf( typeof( Int32 ) ) );
					var p_version = Marshal.AllocHGlobal( Marshal.SizeOf( typeof( Int32 ) ) );

					Interop.zmq_version( p_major, p_minor, p_version );

					m_version = Marshal.ReadInt32( p_major ) +
						"." + Marshal.ReadInt32( p_minor ) +
						"." + Marshal.ReadInt32( p_version );

					Marshal.FreeHGlobal( p_major );
					Marshal.FreeHGlobal( p_minor );
					Marshal.FreeHGlobal( p_version );
				}

				return m_version;
			}
		}

		protected Int64 GetIntSocketOpt( Int32 option )
		{
			var p_value = Marshal.AllocHGlobal( Marshal.SizeOf( typeof( Int64 ) ) );
			var p_length = Marshal.AllocHGlobal( IntPtr.Size );

			if( Interop.zmq_getsockopt( this.Socket, option, p_value, p_length ) != 0 )
			{
				Marshal.FreeHGlobal( p_value );
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			var result = Marshal.ReadInt64( p_value );
			Marshal.FreeHGlobal( p_value );
			Marshal.FreeHGlobal( p_length );
			return result;
		}

		protected void SetIntSocketOpt( Int32 option, Int64 value )
		{
			var p_value = Marshal.AllocHGlobal( Marshal.SizeOf( typeof( Int64 ) ) );

			Marshal.WriteInt64( p_value, value );

			if( Interop.zmq_setsockopt( this.Socket, option, p_value, Marshal.SizeOf( typeof( Int64 ) ) ) != 0 )
			{
				Marshal.FreeHGlobal( p_value );
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			Marshal.FreeHGlobal( p_value );
		}

		protected Byte[] GetByteScoketOpt( Int32 option )
		{
			var p_value = Marshal.AllocHGlobal( Marshal.SizeOf( typeof( Byte ) ) * 4000 );
			var p_length = Marshal.AllocHGlobal( IntPtr.Size );

			if( Interop.zmq_getsockopt( this.Socket, option, p_value, p_length ) != 0 )
			{
				Marshal.FreeHGlobal( p_value );
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			var length_value = Marshal.ReadIntPtr( p_length ).ToInt32();
			Marshal.FreeHGlobal( p_length );

			if( length_value < 4000 )
			{
				Marshal.FreeHGlobal( p_value );
				p_value = Marshal.AllocHGlobal( Marshal.SizeOf( typeof( Byte ) ) * length_value );

				if( Interop.zmq_getsockopt( this.Socket, option, p_value, p_length ) != 0 )
				{
					Marshal.FreeHGlobal( p_value );
					throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
				}

				length_value = Marshal.ReadInt32( p_length );
			}

			var result = new Byte[length_value];
			for( var i = 0; i < result.Length; i++ )
			{
				result[i] = Marshal.ReadByte( p_value, i );
			}

			Marshal.FreeHGlobal( p_value );

			return result;
		}

		protected void SetByteSocketOpt( Int32 option, Byte[] value )
		{
			var p_value = Marshal.AllocHGlobal( Marshal.SizeOf( typeof( Byte ) ) * value.Length );

			for( var i = 0; i < value.Length; i++ )
			{
				Marshal.WriteByte( p_value, i, value[i] );
			}

			if( Interop.zmq_setsockopt( this.Socket, option, p_value, value.Length ) != 0 )
			{
				Marshal.FreeHGlobal( p_value );
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			Marshal.FreeHGlobal( p_value );
		}

		public Int64 HighWatermark { get { return this.GetIntSocketOpt( Interop.ZMQ_HWM ); } set { this.SetIntSocketOpt( Interop.ZMQ_HWM, value ); } }

		public Int64 Swap { get { return this.GetIntSocketOpt( Interop.ZMQ_SWAP ); } set { this.SetIntSocketOpt( Interop.ZMQ_SWAP, value ); } }

		public Int64 Affinity { get { return this.GetIntSocketOpt( Interop.ZMQ_AFFINITY ); } set { this.SetIntSocketOpt( Interop.ZMQ_AFFINITY, value ); } }

		public Byte[] Identity
		{
			get { return this.GetByteScoketOpt( Interop.ZMQ_IDENTITY ); }
			set
			{
				if( value.Length > 255 )
				{
					throw new ArgumentOutOfRangeException( "value", "Identity can not be longer than 255 bytes." );
				}

				this.SetByteSocketOpt( Interop.ZMQ_IDENTITY, value );
			}
		}

		public Int64 MulticastRate { get { return this.GetIntSocketOpt( Interop.ZMQ_RATE ); } set { this.SetIntSocketOpt( Interop.ZMQ_RATE, value ); } }

		public Int64 RecoverySeconds { get { return this.GetIntSocketOpt( Interop.ZMQ_RECOVERY_IVL ); } set { this.SetIntSocketOpt( Interop.ZMQ_RECOVERY_IVL, value ); } }

		public Boolean MulticastLoopback { get { return this.GetIntSocketOpt( Interop.ZMQ_MCAST_LOOP ) > 0; } set { this.SetIntSocketOpt( Interop.ZMQ_MCAST_LOOP, value ? 1 : 0 ); } }

		public Int64 SendBufferSize { get { return this.GetIntSocketOpt( Interop.ZMQ_SNDBUF ); } set { this.SetIntSocketOpt( Interop.ZMQ_SNDBUF, value ); } }

		public Int64 ReceiveBufferSize { get { return this.GetIntSocketOpt( Interop.ZMQ_RCVBUF ); } set { this.SetIntSocketOpt( Interop.ZMQ_RCVBUF, value ); } }

		public Boolean ReceiveMore { get { return this.GetIntSocketOpt( Interop.ZMQ_RCVMORE ) > 0; } }

		public IntPtr FileDescriptor
		{
			get
			{
				var p_value = Marshal.AllocHGlobal( IntPtr.Size );
				var p_length = Marshal.AllocHGlobal( IntPtr.Size );

				if( Interop.zmq_getsockopt( this.Socket, Interop.ZMQ_FD, p_value, p_length ) != 0 )
				{
					Marshal.FreeHGlobal( p_value );
					throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
				}

				var result = Marshal.ReadIntPtr( p_value );
				Marshal.FreeHGlobal( p_value );
				Marshal.FreeHGlobal( p_length );
				return result;
			}
		}

		public Int64 EventState { get { return this.GetIntSocketOpt( Interop.ZMQ_EVENTS ); } }

		public Int64 Type { get { return this.GetIntSocketOpt( Interop.ZMQ_EVENTS ); } }

		public Int64 Linger { get { return this.GetIntSocketOpt( Interop.ZMQ_LINGER ); } set { this.SetIntSocketOpt( Interop.ZMQ_LINGER, value ); } }

		public Int64 ReconnectInterval { get { return this.GetIntSocketOpt( Interop.ZMQ_RECONNECT_IVL ); } set { this.SetIntSocketOpt( Interop.ZMQ_RECONNECT_IVL, value ); } }

		public Int64 ReconnectIntervalMax { get { return this.GetIntSocketOpt( Interop.ZMQ_RECONNECT_IVL_MAX ); } set { this.SetIntSocketOpt( Interop.ZMQ_RECONNECT_IVL_MAX, value ); } }

		public Int64 RecoveryInterval { get { return this.GetIntSocketOpt( Interop.ZMQ_RECOVERY_IVL_MSEC ); } set { this.SetIntSocketOpt( Interop.ZMQ_RECOVERY_IVL_MSEC, value ); } }

		public Int64 MessageBacklog { get { return this.GetIntSocketOpt( Interop.ZMQ_BACKLOG ); } set { this.SetIntSocketOpt( Interop.ZMQ_BACKLOG, value ); } }

		protected void Subscribe( Byte[] prefix )
		{
			this.SetByteSocketOpt( Interop.ZMQ_SUBSCRIBE, prefix );
		}

		protected void Subscribe( String prefix )
		{
			var bytes = Encoding.UTF8.GetBytes( prefix );
			this.Subscribe( bytes );
		}

		protected void Unsubscribe( Byte[] prefix )
		{
			this.SetByteSocketOpt( Interop.ZMQ_UNSUBSCRIBE, prefix );
		}

		protected void Unsubscribe( String prefix )
		{
			var bytes = Encoding.UTF8.GetBytes( prefix );
			this.Unsubscribe( bytes );
		}

		protected void Bind( String address )
		{
			if( Interop.zmq_bind( this.Socket, address ) != 0 )
			{
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}
		}

		protected void Connect( String address )
		{
			if( Interop.zmq_connect( this.Socket, address ) != 0 )
			{
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}
		}

		protected Boolean Send( Byte[] data )
		{
			return this.Send( data, false, false );
		}

		protected Boolean Send( String data )
		{
			return this.Send( Encoding.UTF8.GetBytes( data ) );
		}

		protected Boolean Send( Byte[] data, Boolean nonblocking )
		{
			return this.Send( data, nonblocking, false );
		}

		protected Boolean Send( String data, Boolean nonblocking )
		{
			return this.Send( Encoding.UTF8.GetBytes( data ), nonblocking );
		}

		protected Boolean Send( Byte[] data, Boolean nonblocking, Boolean multipart )
		{
			if( Interop.zmq_msg_init_size( m_send_msg, data.Length ) != 0 )
			{
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			var flags = 0;
			flags |= nonblocking ? Interop.ZMQ_NOBLOCK : 0;
			flags |= multipart ? Interop.ZMQ_SNDMORE : 0;

			Marshal.Copy( data, 0, Interop.zmq_msg_data( m_send_msg ), data.Length );

			var result = Interop.zmq_send( this.Socket, m_send_msg, flags );

			if( Interop.zmq_msg_close( m_send_msg ) != 0 )
			{
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			if( result != 0 )
			{
				if( Interop.zmq_errno() == Interop.EAGAIN )
				{
					return false;
				}

				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			if( this.OnSend != null )
			{
				ThreadPool.QueueUserWorkItem( s => this.OnSend( data, multipart ), null );
			}

			return true;
		}

		protected Boolean Send( String data, Boolean nonblocking, Boolean multipart )
		{
			return this.Send( Encoding.UTF8.GetBytes( data ), nonblocking, multipart );
		}

		protected Boolean Receive( out Byte[] data )
		{
			return this.Receive( out data, false );
		}

		protected Boolean Receive( out String data )
		{
			Byte[] bytes;
			var result = this.Receive( out bytes );
			data = Encoding.UTF8.GetString( bytes );
			return result;
		}

		protected Boolean Receive( out Byte[] data, Boolean nonblocking )
		{
			if( Interop.zmq_msg_init( m_receive_msg ) != 0 )
			{
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			var flags = 0;
			flags |= nonblocking ? Interop.ZMQ_NOBLOCK : 0;

			var result = Interop.zmq_recv( this.Socket, m_receive_msg, flags );

			if( result != 0 )
			{
				data = new Byte[0];

				if( Interop.zmq_errno() == Interop.EAGAIN )
				{
					return false;
				}

				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			var msg_size = Interop.zmq_msg_size( m_receive_msg );
			var p_data = Interop.zmq_msg_data( m_receive_msg );

			if( p_data == IntPtr.Zero )
			{
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			data = new Byte[msg_size];

			Marshal.Copy( p_data, data, 0, data.Length );

			if( Interop.zmq_msg_close( m_receive_msg ) != 0 )
			{
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			return true;
		}

		protected Boolean Receive( out String data, Boolean nonblocking )
		{
			Byte[] bytes;
			var result = this.Receive( out bytes, nonblocking );
			data = Encoding.UTF8.GetString( bytes );
			return result;
		}

		protected event ZmqSendEvent OnSend;

		protected event ZmqReceiveEvent OnReceive;

		private static void Schedule()
		{
			var iterator = s_socket_list.GetEnumerator();

			while( true )
			{
				lock( s_socket_list )
				{
					try
					{
						if( iterator.Current == null || iterator.MoveNext() == false )
						{
							iterator = s_socket_list.GetEnumerator();
							iterator.MoveNext();
						}
					}
					catch( InvalidOperationException )
					{
						iterator = s_socket_list.GetEnumerator();
						iterator.MoveNext();
					}

					if( iterator.Current != null )
					{
						iterator.Current.Poll();
					}
				}

				Thread.Sleep( 1 ); // yield to other threads
			}
		}

		protected virtual void Poll()
		{
			var items = new[]
			{
				new Interop.zmq_pollitem_t
				{
					socket = this.Socket,
					fd = IntPtr.Zero,
					events = Interop.ZMQ_POLLIN,
					revents = 0
				}
			};

			var poll_result = Interop.zmq_poll( items, 1, 10 );

			if( poll_result == -1 )
			{
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			if( poll_result > 0 && ( items[0].revents & Interop.ZMQ_POLLIN ) == Interop.ZMQ_POLLIN && this.OnReceive != null )
			{
				this.OnReceive();
			}
		}

		public virtual void Dispose( Boolean disposing )
		{
			if( disposing )
			{
				
			}

			lock( s_socket_list ) // note- expensive to raipdly create or destroy sockets
			{
				s_socket_list.Remove( this );
			}

			Marshal.FreeHGlobal( m_send_msg );
			Marshal.FreeHGlobal( m_receive_msg );

			if( Interop.zmq_close( this.Socket ) != 0 ) {
				throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
			}

			if( s_socket_list.Count < 1 )
			{
				s_poller.Abort();
				s_poller.Join();

				s_poller = null;

				s_socket_list = null;

				if( Interop.zmq_term( s_context ) != 0 )
				{
					throw new ZmqException( Interop.zmq_errno(), Marshal.PtrToStringAnsi( Interop.zmq_strerror( Interop.zmq_errno() ) ) );
				}

				s_context = IntPtr.Zero;
			}
		}

		public virtual void Dispose()
		{
			this.Dispose( true );
			GC.SuppressFinalize( this );
		}

		~ZmqSocket()
		{
			this.Dispose( false );
		}
	}
}
